from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.messages import TextMessage
from autogen_agentchat.ui import Console
from autogen_core import CancellationToken
from ..model import model_client

streaming_assistant = AssistantAgent(
    name="assistant",
    model_client=model_client,
    system_message="You are a helpful assistant.",
    model_client_stream=True,  # Enable streaming tokens.
)

async def main():
    # Use an async function and asyncio.run() in a script.
    async for message in streaming_assistant.on_messages_stream(  # type: ignore
        [TextMessage(content="Name two cities in South America", source="user")],
        cancellation_token=CancellationToken(),
    ):
        print(message)